import numpy as np
import numpy.ctypeslib as npC
import ctypes as C
import os
from scipy.spatial import KDTree

_path = os.path.dirname(__file__)
_lib = npC.load_library("librptree.so",_path)

## RPTree

#void *RPTree_new( double *data,
#                  unsigned int rows,
#                  unsigned int columns,
#                  unsigned int min_size,
#                  unsigned int maxdepth,
#                  double c,
#                  unsigned int *leavesContent );
_lib.RPTree_new.restype = C.c_void_p
_lib.RPTree_new.argtypes = [
	npC.ndpointer(ndim = 2, dtype = 'float64', flags =  'c_contiguous'),
	C.c_uint,
        C.c_uint,
	C.c_uint,
	C.c_uint,
        C.c_double,
	npC.ndpointer(ndim = 1, dtype = 'uint32', flags =  'c_contiguous')
]

#void *RPTree_load( const char *file );
_lib.RPTree_load.restype = C.c_void_p
_lib.RPTree_load.argtypes = [ C.c_char_p ]

#void RPTree_save( void *rptree, const char *file );
_lib.RPTree_save.restype = None
_lib.RPTree_save.argtypes = [
        C.c_void_p,
        C.c_char_p
]

#void RPTree_print( void *rptree );
_lib.RPTree_print.restype = None
_lib.RPTree_print.argtypes = [ C.c_void_p ]

#void RPTree_delete( void *rptree );
_lib.RPTree_delete.restype = None
_lib.RPTree_delete.argtypes = [ C.c_void_p ]

#unsigned int RPTree_use( void *rptree, double *data );
_lib.RPTree_use.restype = C.c_uint
_lib.RPTree_use.argtypes = [
        C.c_void_p,
        npC.ndpointer(ndim = 1, dtype = 'float64', flags =  'c_contiguous')
]

#void RPTree_multi_use( void *rptree, double *data, unsigned int rows, unsigned int *result );
_lib.RPTree_multi_use.restype = None
_lib.RPTree_multi_use.argtypes = [
        C.c_void_p,
        npC.ndpointer(ndim = 2, dtype = 'float64', flags =  'c_contiguous'),
        C.c_uint, npC.ndpointer(ndim = 1, dtype = 'uint32', flags =  'c_contiguous')
]

#void RPTree_get_v( void *rptree, double *v );
_lib.RPTree_get_v.restype = None
_lib.RPTree_get_v.argtypes = [
        C.c_void_p,
        npC.ndpointer(ndim=1, dtype = 'float64', flags =  'c_contiguous')
]

#void *RPTree_get_left( void *rptree );
_lib.RPTree_get_left.restype = C.c_void_p
_lib.RPTree_get_left.argtypes = [ C.c_void_p ]

#void *RPTree_get_right( void *rptree );
_lib.RPTree_get_right.restype = C.c_void_p
_lib.RPTree_get_right.argtypes = [ C.c_void_p ]

#int RPTree_get_usedRule( void *rptree );
_lib.RPTree_get_usedRule.restype = C.c_int
_lib.RPTree_get_usedRule.argtypes = [ C.c_void_p ]

#double RPTree_get_threshold( void *rptree );
_lib.RPTree_get_threshold.restype = C.c_double
_lib.RPTree_get_threshold.argtypes = [ C.c_void_p ]

#double RPTree_get_mean( void *rptree );
_lib.RPTree_get_mean.restype = C.c_double
_lib.RPTree_get_mean.argtypes = [ C.c_void_p ]

#unsigned int RPTree_get_dim( void *rptree );
_lib.RPTree_get_dim.restype = C.c_uint
_lib.RPTree_get_dim.argtypes = [ C.c_void_p ]

#unsigned int RPTree_get_depth( void *rptree );
_lib.RPTree_get_depth.restype = C.c_uint
_lib.RPTree_get_depth.argtypes = [ C.c_void_p ]

#void *RPTree_get_leaf_by_id( void *rptree, unsigned int id );
_lib.RPTree_get_leaf_by_id.restype = C.c_void_p
_lib.RPTree_get_leaf_by_id.argtypes = [ C.c_uint ]

#bool RPTRee_isLeaf(void* rptree);
_lib.RPTree_isLeaf.retype = C.c_bool
_lib.RPTree_isLeaf.argtypes = [ C.c_void_p ]

## RPLeaf

#void *RPLeaf_new( unsigned int depth, unsigned int id );
_lib.RPLeaf_new.restype = C.c_void_p
_lib.RPLeaf_new.argtypes = [
        C.c_uint,
        C.c_uint
]

#unsigned int RPLeaf_use( void *rpleaf, double *data );
_lib.RPLeaf_use.restype = C.c_uint
_lib.RPLeaf_use.argtypes = [
        C.c_void_p,
        npC.ndpointer(ndim = 1, dtype = 'float64', flags =  'c_contiguous')
]

#void RPLeaf_delete( void *rpleaf );
_lib.RPLeaf_delete.restype = None
_lib.RPLeaf_delete.argtypes = [ C.c_void_p ]

#unsigned int RPLeaf_get_id( void *rpleaf );
_lib.RPLeaf_get_id.restype = C.c_uint
_lib.RPLeaf_get_id.argtypes = [ C.c_void_p ]

RULE_NONE = -1
RULE_HYPERPLANE = 0
RULE_MEDIAN = 1

def node_factory(node_obj):
	if _lib.RPTree_isLeaf(node_obj) :
		return RPLeaf(node_obj)
	else :
		return RPTree(node_obj)

def build_tree(data, min_size, maxdepth, c):
	nrow,ncol = data.shape
	leavesContent = np.zeros(nrow,'uint32')
	node = _lib.RPTree_new( data, nrow, ncol, min_size, maxdepth, c, leavesContent)
	return RPTree(node,is_root = True),leavesContent

def explore_tree(node,function,*args):
	function(node,*args)
	if not node.isLeaf() :
		explore_tree(node.left,function,*args)
		explore_tree(node.right,function,*args)

def get_rules(node):
	rules = []
	def _func(node,rules):
		if not node.isLeaf() : rules.append(node.usedRule)
	explore_tree(node,_func,rules)
	return np.array(rules)

def get_centroids(node):
	centr = []
	def _func(node,centr):
		if node.isLeaf() : centr.append(node.v)
	explore_tree(node,_func,centr)
	return np.array(centr)

def compute_vq_distortion(tree,data):
	centrs = get_centroids(tree)
	kdt = KDTree(centrs)
	d,indx = kdt.query(data)
	return (d**2).sum()

def compute_hyperplane_threshold(v,data):
	a = np.dot(data,v)
	aAcc = a.cumsum()
	count = np.array(range(len(a))) + 1
	mu1 = aAcc[:-1]/count[:-1]
	mu2 = (aAcc[-1] - aAcc[:-1])/count[:-1][::-1]
	
	v1 = ((a[:-1]-mu1)**2).cumsum()
	v2 = ((a[:-1]-mu2)**2).cumsum()
	v2 = v2[-1] - v2[::-1]
	ci = v1+v2
	imin = ci.argmin()
	return (a[imin]+a[imin+1])/2.

def compute_hyperplane_threshold2(v,data):
	a = np.dot(data,v)
	aAcc = a.cumsum()
	aAcc2 = (a**2).cumsum()[:-1]
	count = np.array(range(len(a))) + 1
	mu1 = aAcc[:-1]/count[:-1]
	mu2 = (aAcc[-1] - aAcc[:-1])/count[:-1][::-1]

	ci = aAcc2 - count[:-1]*mu1*mu1 + (aAcc2[-1] - aAcc2) - count[:-1][::-1] * mu2 * mu2
	imin = ci.argmin()
	return (a[imin]+a[imin+1])/2.

def compute_hyperplane_threshold3(v,data):
	a = np.dot(data,v)
	aAcc = a.cumsum()
	aAcc2 = (a**2).cumsum()[:-1]
	count = np.array(range(len(a))) + 1
	mu1 = aAcc[:-1]/count[:-1]
	mu2 = (aAcc[-1] - aAcc[:-1])/count[:-1][::-1]

	ci = aAcc2 - count[:-1]*mu1*mu1 + (aAcc2[-1] - aAcc2) - count[:-1][::-1] * mu2 * mu2
	imin = ci.argmin()
	return (a[imin]+a[imin+1])/2.


class RPTree(object):
        def __init__(self, node_obj, is_root = False ):
		self._object = node_obj
		self.is_root = is_root
		self._dim = None
		self._lib = _lib

	def save(self, file):
			self._lib.RPTree_save( self._object, file )

	def Print(self):
			self._lib.RPTree_print( self._object )

	def __del__(self):
		if self.is_root : self._lib.RPTree_delete( self._object )

	def use(self, data):
		if data.ndim == 1 :
			return self._lib.RPTree_use( self._object, data)
		else :
			ndata = data.shape[0]
			node_ids = np.zeros(ndata,'uint32')
			self._lib.RPTree_multi_use(self._object,data,ndata,node_ids)
			return node_ids

	@property
	def v(self):
		v = np.zeros(self.dim)
		self._lib.RPTree_get_v(self._object, v)
		return v

	@property
	def left(self):
		return node_factory(self._lib.RPTree_get_left(self._object))

	@property
	def right(self):
		return node_factory(self._lib.RPTree_get_right(self._object))

	@property
	def usedRule(self):
		return self._lib.RPTree_get_usedRule(self._object)

	@property
	def threshold(self):
		return self._lib.RPTree_get_threshold(self._object)

	@property
	def mean(self):
		return self._lib.RPTree_get_mean(self._object)

	@property
	def dim(self):
		if self._dim is None :
			self._dim = self._lib.RPTree_get_dim(self._object)
		return self._dim

	def get_leaf_by_id(self, id):
		return node_factory(self._lib.RPTree_get_leaf_by_id(self._object, id))
	
	def isLeaf(self): return False


class RPLeaf(RPTree):
	def __init__(self, node_object):
		RPTree.__init__(self,node_object)

	def use(self, data):
		raise Exception("Shouldn't be called directly")

	def get_id(self):
		return self._lib.RPLeaf_get_id(self._object)
				   
	def isLeaf(self): return True

if __name__ == '__main__' :
	data = np.random.rand(100,1).repeat(3,0)
	tree,leavesContent = build_tree(data,10, 2, 10)
	tree.use(data)
